##sp4  class_sim
from ndemag import setNdemag
import numpy as np
from matplotlib import pyplot as plt
from multiprocessing import Pool
import plotly.express as px
import time
def N(name,n,dx):
    n_demag = np.zeros([2*i-1 for i in n] + [6,])
    results=[]
    p=Pool(6)
    for i, per_fun in enumerate(((0,1,2),(0,1,2),(0,2,1),(1,2,0),(1,2,0),(2,0,1))):
        result_=p.apply_async(func=setNdemag,args=(i,list(per_fun),list(n),list(dx),n_demag ,))
        results.append(result_)
    p.close()
    p.join()
    for i in range(6):
        n_demag[:,:,:,i]=results[i].get()
    np.save(str(name)+"_ndemag"+".npy",n_demag)
if __name__=="__main__":
    st=time.perf_counter()
    name="sp4"
    n=[100,25,1]
    dx=[5e-9,5e-9,3e-9]
    mu0=4e-7 * np.pi
    N(name=name,n=n,dx=dx)
    import class_sim as sim#避免一些multiprocess的错误
    sim.sim(
    name=name,
    namedemag="sp4_ndemag",
    n=n,
    dx=dx,
    m0=[1,0.1,0],
    gamma = 2.211e5,
    ms    = 800e3,
    A     = 13e-12,
    alpha=1,
    h_zee=lambda t:[0,0,0],
    T=1e-9,
    T0=0,
    fixdt=1e-12,).run()

    m0=np.load(str(name)+"SIMnow"+".npy")
    sim.sim(
    name=name,
    namedemag="sp4_ndemag",
    n=[100,25,1],
    dx=[5e-9,5e-9,3e-9],
    m0=m0,
    gamma = 2.211e5,
    ms    = 800e3,
    A     = 13e-12,
    alpha=0.02,
    h_zee=lambda t:[-24.6e-3/mu0, +4.3e-3/mu0, 0.0],
    T=1e-9,
    T0=0,
    # fixdt=1e-12,
    ).run()
    en=time.perf_counter()
    print(en-st)

#======visual===================================================
    Y=np.load(str(name)+"SIM"+".npy")
    R=np.load(str(name)+"SIM"+"error.npy",allow_pickle=True)
    m_avg=Y
    plt.subplot(1, 2, 2)
    plt.plot(m_avg[:,0], m_avg[:,1]) 
    plt.plot(m_avg[:,0], m_avg[:,2])
    plt.plot(m_avg[:,0], m_avg[:,3])

    plt.xlabel("Time (s)")
    plt.ylabel("Magnetization")
    plt.title("m")

    plt.subplot(1, 2, 1)
    plt.plot(range(len(R)),R)

    plt.xlabel("step")
    plt.ylabel("error")
    plt.title("R")

    plt.show()
    input()
